use pest::Parser;
use pest_derive::Parser;
use std::collections::HashMap;
use std::env;
use std::fs;
use std::rc::Rc;

#[derive(Parser)]
#[grammar = "grammar.pest"]
struct LispParser;

#[derive(Clone, Debug)]
enum LispValue {
    Symbol(String),
    Number(f64),
    String(String),
    List(Vec<LispValue>),
    Function(fn(&[LispValue]) -> Result<LispValue, String>),
    Lambda(Vec<String>, Vec<LispValue>, Rc<Environment>),
}

#[derive(Debug, Clone)]
struct Environment {
    values: HashMap<String, LispValue>,
    parent: Option<Rc<Environment>>,
}

impl Environment {
    fn new() -> Self {
        let mut env = Environment {
            values: HashMap::new(),
            parent: None,
        };
        env.define_natives();
        env
    }

    fn with_parent(parent: Rc<Environment>) -> Self {
        let mut env = Environment {
            values: HashMap::new(),
            parent: Some(parent),
        };
        env.define_natives();
        env
    }
    fn define(&mut self, name: String, value: LispValue) {
        // 如果是Lambda，将当前环境注入以支持递归
        match &value {
            LispValue::Lambda(params, body, _) => {
                let mut new_env = Environment::new();
                new_env.parent = Some(Rc::new(self.clone()));
                self.values.insert(
                    name.clone(),
                    LispValue::Lambda(params.clone(), body.clone(), Rc::new(new_env)),
                );
            }
            _ => {
                self.values.insert(name, value);
            }
        }
    }

    fn define_natives(&mut self) {
        self.values.insert(
            "+".to_string(),
            LispValue::Function(|args| match &args[0] {
                LispValue::Number(n) => {
                    let mut result = *n;
                    for arg in &args[1..] {
                        if let LispValue::Number(x) = arg {
                            result += x;
                        }
                    }
                    Ok(LispValue::Number(result))
                }
                LispValue::String(s) => {
                    let mut result = s.clone();
                    for arg in &args[1..] {
                        if let LispValue::String(x) = arg {
                            result.push_str(x);
                        }
                    }
                    Ok(LispValue::String(result))
                }
                _ => Err("Invalid arguments for +".to_string()),
            }),
        );

        self.values.insert(
            "-".to_string(),
            LispValue::Function(|args| {
                if let LispValue::Number(n) = args[0] {
                    if args.len() == 1 {
                        return Ok(LispValue::Number(-n));
                    }
                    let mut result = n;
                    for arg in &args[1..] {
                        if let LispValue::Number(x) = arg {
                            result -= x;
                        }
                    }
                    Ok(LispValue::Number(result))
                } else {
                    Err("Invalid arguments for -".to_string())
                }
            }),
        );

        self.values.insert(
            "*".to_string(),
            LispValue::Function(|args| {
                if let LispValue::Number(n) = args[0] {
                    let mut result = n;
                    for arg in &args[1..] {
                        if let LispValue::Number(x) = arg {
                            result *= x;
                        }
                    }
                    Ok(LispValue::Number(result))
                } else {
                    Err("Invalid arguments for *".to_string())
                }
            }),
        );

        self.values.insert(
            "=".to_string(),
            LispValue::Function(|args| {
                if args.len() != 2 {
                    return Err("= requires exactly 2 arguments".to_string());
                }
                match (&args[0], &args[1]) {
                    (LispValue::Number(a), LispValue::Number(b)) => {
                        Ok(LispValue::Number(if (a - b).abs() < f64::EPSILON {
                            1.0
                        } else {
                            0.0
                        }))
                    }
                    _ => Err("Arguments to = must be numbers".to_string()),
                }
            }),
        );

        // 小于操作符
        self.values.insert(
            "<".to_string(),
            LispValue::Function(|args| {
                if args.len() != 2 {
                    return Err("< requires exactly 2 arguments".to_string());
                }
                match (&args[0], &args[1]) {
                    (LispValue::Number(a), LispValue::Number(b)) => {
                        Ok(LispValue::Number(if a < b { 1.0 } else { 0.0 }))
                    }
                    _ => Err("Arguments to < must be numbers".to_string()),
                }
            }),
        );

        // 大于操作符
        self.values.insert(
            ">".to_string(),
            LispValue::Function(|args| {
                if args.len() != 2 {
                    return Err("> requires exactly 2 arguments".to_string());
                }
                match (&args[0], &args[1]) {
                    (LispValue::Number(a), LispValue::Number(b)) => {
                        Ok(LispValue::Number(if a > b { 1.0 } else { 0.0 }))
                    }
                    _ => Err("Arguments to > must be numbers".to_string()),
                }
            }),
        );
    }

    fn get(&self, name: &str) -> Option<LispValue> {
        match self.values.get(name) {
            Some(value) => Some(value.clone()),
            None => {
                if let Some(parent) = &self.parent {
                    parent.get(name)
                } else {
                    None
                }
            }
        }
    }

    fn set(&mut self, name: String, value: LispValue) {
        self.values.insert(name, value);
    }
}

fn parse_expr(pair: pest::iterators::Pair<Rule>) -> LispValue {
    match pair.as_rule() {
        Rule::number => LispValue::Number(pair.as_str().parse().unwrap()),
        Rule::string => {
            let s = pair.as_str();
            LispValue::String(s[1..s.len() - 1].to_string())
        }
        Rule::symbol => LispValue::Symbol(pair.as_str().to_string()),
        Rule::list => {
            let exprs: Vec<LispValue> = pair.into_inner().map(parse_expr).collect();
            LispValue::List(exprs)
        }
        Rule::quote => {
            let expr = parse_expr(pair.into_inner().next().unwrap());
            LispValue::List(vec![LispValue::Symbol("quote".to_string()), expr])
        }
        Rule::expr => parse_expr(pair.into_inner().next().unwrap()),
        _ => unreachable!(),
    }
}

fn eval(expr: &LispValue, env: &mut Environment) -> Result<LispValue, String> {
    match expr {
        LispValue::Symbol(s) => match env.get(s) {
            Some(val) => Ok(val),
            None => Err(format!("Undefined symbol: {}", s)),
        },
        LispValue::Number(_) | LispValue::String(_) => Ok(expr.clone()),
        LispValue::List(list) => {
            if list.is_empty() {
                return Ok(LispValue::List(vec![]));
            }

            match &list[0] {
                LispValue::Symbol(s) => match s.as_str() {
                    "define" => {
                        if list.len() != 3 {
                            return Err("define requires exactly 2 arguments".to_string());
                        }
                        if let LispValue::Symbol(name) = &list[1] {
                            let value = eval(&list[2], env)?;
                            env.define(name.clone(), value.clone());
                            Ok(value)
                        } else {
                            Err("First argument to define must be a symbol".to_string())
                        }
                    }
                    "quote" => {
                        if list.len() != 2 {
                            return Err("quote requires exactly 1 argument".to_string());
                        }
                        Ok(list[1].clone())
                    }
                    "lambda" => {
                        if list.len() < 3 {
                            return Err("lambda requires at least 2 arguments".to_string());
                        }
                        if let LispValue::List(params) = &list[1] {
                            let mut param_names = Vec::new();
                            for param in params {
                                if let LispValue::Symbol(name) = param {
                                    param_names.push(name.clone());
                                } else {
                                    return Err("Lambda parameters must be symbols".to_string());
                                }
                            }
                            let body = list[2..].to_vec();
                            Ok(LispValue::Lambda(param_names, body, Rc::new(env.clone())))
                        } else {
                            Err("Lambda parameters must be a list".to_string())
                        }
                    }
                    "if" => {
                        if list.len() != 4 {
                            return Err("if requires exactly 3 arguments".to_string());
                        }
                        let condition = eval(&list[1], env)?;
                        match condition {
                            LispValue::Number(n) => {
                                if n != 0.0 {
                                    eval(&list[2], env)
                                } else {
                                    eval(&list[3], env)
                                }
                            }
                            _ => Err("Condition must evaluate to a number".to_string()),
                        }
                    }
                    _ => {
                        let func = eval(&list[0], env)?;
                        let mut args = Vec::new();
                        for arg in &list[1..] {
                            args.push(eval(arg, env)?);
                        }
                        match func {
                            LispValue::Function(f) => f(&args),
                            LispValue::Lambda(params, body, closure_env) => {
                                let mut new_env = Environment::with_parent(closure_env);
                                for (param, arg) in params.iter().zip(args) {
                                    new_env.set(param.clone(), arg);
                                }
                                let mut result = LispValue::List(vec![]);
                                for expr in body {
                                    result = eval(&expr, &mut new_env)?;
                                }
                                Ok(result)
                            }
                            _ => Err("First element of list must be a function".to_string()),
                        }
                    }
                },
                _ => Err("First element of list must be a symbol".to_string()),
            }
        }
        LispValue::Function(_) => Ok(expr.clone()),
        LispValue::Lambda(_, _, _) => Ok(expr.clone()),
    }
}

fn print_value(value: &LispValue) -> String {
    match value {
        LispValue::Symbol(s) => s.clone(),
        LispValue::Number(n) => n.to_string(),
        LispValue::String(s) => format!("\"{}\"", s),
        LispValue::List(list) => {
            let items: Vec<String> = list.iter().map(print_value).collect();
            format!("({})", items.join(" "))
        }
        LispValue::Function(_) => "#<function>".to_string(),
        LispValue::Lambda(_, _, _) => "#<lambda>".to_string(),
    }
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let args: Vec<String> = env::args().collect();
    if args.len() != 2 {
        return Err("Usage: mylisp <filename>".into());
    }

    let filename = &args[1];
    let source = fs::read_to_string(filename)?;

    let mut env = Environment::new();
    let pairs = LispParser::parse(Rule::program, &source)?;

    for pair in pairs {
        match pair.as_rule() {
            Rule::program => {
                for expr_pair in pair.into_inner() {
                    if expr_pair.as_rule() == Rule::expr {
                        let expr = parse_expr(expr_pair);
                        let result = eval(&expr, &mut env)?;
                        println!("{}", print_value(&result));
                    }
                }
            }
            _ => unreachable!(),
        }
    }

    Ok(())
}
